#include "../MulticastUdpSocket.h"

#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>

// TODO: move this to general part
#include <string.h>

namespace Greatewall
{
    namespace Network
    {
        SysSockHandler
        MulticastUdpSocket::createMcUdpSock()
        {
            return socket(AF_INET, SOCK_DGRAM, 0);
        }

        bool
        MulticastUdpSocket::send(const char* buf, ssize_t& size, bool non_block)
        {
            if (_mode == RECV)
            {
                // TODO: Error handling
                return false;
            }

            if (!_inited)
            {
                // TODO: Error handling
                return false;
            }

            if (!_joined)
            {
                // TODO: Error handling
                return false;
            }

            return IpSocket::send(buf, size, non_block);
        }

        bool
        MulticastUdpSocket::recv(char*& buf, ssize_t& size, bool non_block)
        {
            if (_mode == SEND)
            {
                // TODO: Error handling
                return false;
            }

            if (!_inited)
            {
                // TODO: Error handling
                return false;
            }

            if (!_joined)
            {
                // TODO: Error handling
                return false;
            }

            return IpSocket::recv(buf, size, non_block);
        }

        bool
        MulticastUdpSocket::initMcUdpSock(std::string& mc, short port, std::string& netif, MCMode mode)
        {
            if (_inited)
            {
                // Cannot initialize twice
                // TODO: Error handling
                return false;
            }

            if (_hdl < 0)
            {
                /// Create the socket if there is not
                _hdl = createMcUdpSock();

                /// TODO: Error handling
                if (_hdl < 0)
                {
                    return false;
                }
            }

            switch (mode)
            {
                case RECV:
                {
                    int         flag_on = 1;

                    /// Make port reuseble
                    if ((setsockopt(_hdl, SOL_SOCKET, SO_REUSEADDR, &flag_on, sizeof(flag_on))) < 0)
                    {
                        /// TODO: Error handling
                        return false;
                    }

                    /// Construct a multicast incoming address struct
                    memset(&_in, 0, sizeof(_in));
                    _in.sin_family     = AF_INET;
                    _in.sin_addr.s_addr= htonl(INADDR_ANY);
                    _in.sin_port       = htons(port);

                    /// Bind the socket to incoming address
                    if ((bind(_hdl, (sockaddr*)&_in, sizeof(_in))) < 0)
                    {
                        /// TODO: Error handling
                        return false;
                    }

                    ///Check the validity of MC address
                    in_addr_t mc_addr = 0;;
                    if (!inet_aton(mc.c_str(), (in_addr*)&mc_addr) || (mc_addr & 0xF000) != 0xE000)
                    {
                        /// TODO: Error handling
                        return false;
                    }

                    break;
                }
                case SEND:
                {
                    unsigned char mc_ttl = 1;

                    /// Set the sock option
                    if ((setsockopt(_hdl, IPPROTO_IP, IP_MULTICAST_TTL, (void*)&mc_ttl, sizeof(mc_ttl))) < 0)
                    {
                        /// TODO: Error handling
                        return false;
                    }

                    /// Check MC address, and construct out address
                    in_addr_t mc_addr = 0; inet_aton(mc.c_str(), (in_addr*)&mc_addr);
                    if ((mc_addr & 0xF000) != 0xE000)
                    {
                        /// TODO: Error handling
                        return false;
                    }
                    memset(&_out, 0, sizeof(_out));
                    _out.sin_family     = AF_INET;
                    _out.sin_addr.s_addr= mc_addr;
                    _out.sin_port       = htons(port);

                    break;
                }
                default:
                {
                    // WRONG MODE
                    // TODO: set ERROR
                    return false;
                }
            }

            _mc     = mc;
            _port   = port;
            _netif  = netif;
            _mode   = mode;
            _inited = true;
            return true;
        }

        bool
        MulticastUdpSocket::join()
        {
            if (!_inited)
            {
                // TODO: Error handling
                // Must called after initialization
                return false;
            }

            if (_joined)
            {
                // TODO: Error handling
                return false;
            }

            if (_hdl < 0)
            {
                /// TODO: Error handling
                return false;
            }

            ip_mreq     mc_req;

            ///Check the validity of MC address
            in_addr_t mc_addr = 0;;
            if (!inet_aton(_mc.c_str(), (in_addr*)&mc_addr) || (mc_addr & 0xF000) != 0xE000)
            {
                /// TODO: Error handling
                return false;
            }

            in_addr_t if_addr = htonl(INADDR_ANY);
            if (!_netif.empty() && !inet_aton(_netif.c_str(), (in_addr*)&if_addr))
            {
                /// TODO: Error handling
                return false;
            }

            /// construct IGMP join request
            memset(&mc_req, 0, sizeof(mc_req));
            mc_req.imr_multiaddr.s_addr = mc_addr;
            mc_req.imr_interface.s_addr = if_addr;

            /// join the multicast channel
            if ((setsockopt(_hdl, IPPROTO_IP, IP_ADD_MEMBERSHIP, (void*)&mc_req, sizeof(mc_req))) < 0)
            {
                /// TODO: Error handling
                return false;
            }

            _joined = true;

            return true;
        }

        bool
        MulticastUdpSocket::drop()
        {
            if (!_inited)
            {
                // TODO: Error handling
                return false;
            }

            if (!_joined)
            {
                // TODO: Error handling
                return false;
            }

            ip_mreq     mc_req;

            ///Check the validity of MC address:
            in_addr_t mc_addr = 0;;
            if (!inet_aton(_mc.c_str(), (in_addr*)&mc_addr) || (mc_addr & 0xF000) != 0xE000)
            {
                /// TODO: Error handling
                return false;
            }

            in_addr_t if_addr = htonl(INADDR_ANY);
            if (!_netif.empty() && !inet_aton(_netif.c_str(), (in_addr*)&if_addr))
            {
                /// TODO: Error handling
                return false;
            }

            /// construct IGMP join request
            memset(&mc_req, 0, sizeof(mc_req));
            mc_req.imr_multiaddr.s_addr = mc_addr;
            mc_req.imr_interface.s_addr = if_addr;

            /// drop the multicast channel
            if ((setsockopt(_hdl, IPPROTO_IP, IP_DROP_MEMBERSHIP, (void*)&mc_req, sizeof(mc_req))) < 0)
            {
                /// TODO: Error handling
                return false;
            }

            _joined = false;

            return true;
        }

        void
        MulticastUdpSocket::close()
        {
            if (_hdl > 0)
            {
                ::close(_hdl);
                _hdl = -1;
                _inited = false;
            }
        }
    }
}

